/*
 * IORNN.cpp
 *
 *  Created on: Dec 17, 2014
 *      Author: Phong
 */

#include "IORNN.h"
#include "Utils.h"
#include <fstream>

using namespace std;

IORNN::IORNN(int dim, Dictionary* vocaDic, Dictionary* posDic, Dictionary* deprelDic,
		vector<Matrix*>& L, int cmplInside) {
	this->vocaDic = vocaDic->share();
	this->posDic = posDic->share();
	this->deprelDic = deprelDic->share();

	this->dim = dim;
	this->wdim = L[0]->length;
	this->params = NULL;
	this->initParam(L);

	this->completeInside = cmplInside;
}

IORNN::~IORNN() {
	Dictionary::free(vocaDic);
	Dictionary::free(posDic);
	Dictionary::free(deprelDic);
	Param::free(params);
}

void IORNN::initParam(vector<Matrix*>& L) {
	params = new Param(dim, wdim, vocaDic->size(), posDic->size(), deprelDic->size(), largeRandRange, smallRandRange);
	for (int i = 0; i < params->L.size(); i++)
		params->L[i]->copy(L[i]);
}

void IORNN::save(string fname) {
	ofstream f(fname);

	f << dim << " " << wdim << " " << completeInside << endl;

	vocaDic->save(f);
	posDic->save(f);
	deprelDic->save(f);
	params->save(f);

	f << endl;
	f.close();
}

IORNN* IORNN::load(string fname) {
	ifstream f(fname);
	IORNN* net = new IORNN();

	f >> net->dim;
	f >> net->wdim;
	f >> net->completeInside;

	net->vocaDic = Dictionary::load(f);
	net->posDic = Dictionary::load(f);
	net->deprelDic = Dictionary::load(f);
	net->params = Param::load(f);

	f.close();
	return net;
}

Param* IORNN::createGrad() {
	Param* p = new Param(dim, wdim, vocaDic->size(), posDic->size(), deprelDic->size(), 0, 0);
	p->fill(0);
	return p;
}

Matrix* IORNN::func(Matrix* X) {
	Matrix* F = new Matrix(X->rows, X->cols);
	for (int i = 0; i < F->length; i++) {
		F->data[i] = X->data[i] / (1 + abs(X->data[i]));
	}
	return F;
}

Matrix* IORNN::funcPrime(Matrix* F) {
	Matrix* Fp = new Matrix(F->rows, F->cols);
	for (int i = 0; i < F->length; i++) {
		Fp->data[i] = 1 - abs(F->data[i]);
		Fp->data[i] = Fp->data[i] * Fp->data[i];
	}
	return Fp;
}

Matrix* IORNN::sigmoid(Matrix* X) {
	Matrix* F = new Matrix(X->rows, X->cols);
	for (int i = 0; i < F->length; i++) {
		F->data[i] = 1 / (1 + exp(-X->data[i]));
	}
	return F;
}

void IORNN::forwardInside(Deptree* tree) {
	// computing inner for heads
	tree->headInner = new Matrix*[tree->nNodes];

	Matrix* input = new Matrix(dim);
	for (int i = 1; i < tree->nNodes; i++) {
		Matrix* Li = params->L[tree->word[i]];
		Matrix::gemv(1, CblasNoTrans, params->Ww, Li, 0, input);
		input->addi(params->Lpos[tree->pos[i]])
            		->addi(params->Lcap[tree->cap[i]])
					->addi(params->bw);
		tree->headInner[i] = func(input);
	}
	tree->headInner[0] = params->rootInner->dup();

	// compute phrase inner
	if (completeInside != CMPL_NONE) {
		tree->phraseInner = new Matrix*[tree->nNodes];
		for (int i = tree->nNodes-1; i >= 0; i--) {
			int nChildren = tree->nChildren[i][0] + tree->nChildren[i][1];
			if (nChildren == 0)
				tree->phraseInner[i] = tree->headInner[i]->dup();
			else {
				input->fill(0);
				for (int dir = 0; dir < 2; dir++) {
					for (int j = 0; j < tree->nChildren[i][dir]; j++) {
						int child = tree->children[i][dir][j];
						Matrix* W = params->Wi[dir][tree->deprel[child]];
						Matrix::gemv(1, CblasNoTrans, W, tree->phraseInner[child], 1, input);
					}
				}
				Matrix::gemv(1, CblasNoTrans, params->Wih, tree->headInner[i], 1./nChildren, input);
				input->addi(params->bi);
				tree->phraseInner[i] = func(input);
			}
		}
	}
	delete input;
}

real IORNN::forwardOutside(Deptree* tree) {
	tree->outer = new Matrix*[tree->nNodes]; // full outer
	tree->pOuter = new Matrix*[tree->nNodes]; // partial outer
	tree->EOCOuter = new Matrix**[2] {
		new Matrix*[tree->nNodes],
		new Matrix*[tree->nNodes] };

	// precompute
	Matrix*** piProds = new Matrix**[2]{new Matrix*[tree->nNodes],
		new Matrix*[tree->nNodes]}; // Wo * tree->phraseInner
	Matrix*** hiProds = new Matrix**[2]{new Matrix*[tree->nNodes],
		new Matrix*[tree->nNodes]}; // Wo * tree->headInner
	for (int i = 1; i < tree->nNodes; i++) {
		int dir = tree->dir[i];
		int deprel = tree->deprel[i];
		Matrix* W = params->Wo[dir][deprel];

		hiProds[dir][i] = new Matrix(dim);
		Matrix::gemv(1, CblasNoTrans, W, tree->headInner[i], 0, hiProds[dir][i]);
		if (completeInside != CMPL_NONE) {
			piProds[dir][i] = new Matrix(dim);
			Matrix::gemv(1, CblasNoTrans, W, tree->phraseInner[i], 0, piProds[dir][i]);
		}
	}

	// compute partial/full outers
	Matrix* inputParent = new Matrix(dim);
	Matrix* input = new Matrix(dim);
	Matrix* inputHead = new Matrix(dim);

	for (int i = 0; i < tree->nNodes; i++) {
		// compute full outer
		if (i == 0) {// ROOT
			tree->outer[i] = params->anonOuter->dup();
			tree->pOuter[i] = Matrix::zeros(dim);
		}
		else {
			int parent = tree->parent[i];
			int dir = tree->dir[i];
			Matrix::gemv(1, CblasNoTrans, params->Woh[dir], tree->headInner[parent], 0, inputParent);
			Matrix::gemv(1, CblasNoTrans, params->Wop[dir], tree->outer[parent], 1, inputParent);
			inputParent->addi(params->bo[dir]);

			int nSisters = tree->nChildren[parent][0] + tree->nChildren[parent][1] - 1;
			if (nSisters == 0) {
				inputParent->addi(params->anonInner[dir]);
				tree->outer[i] = func(inputParent);
			}
			else {
				input->fill(0);
				for (int sdir = 0; sdir < 2; sdir++) {
					for (int j = 0; j < tree->nChildren[parent][sdir]; j++) {
						int sister = tree->children[parent][sdir][j];
						if (sister < i) {
							if (completeInside == CMPL_LEFT2RIGHT)
								input->addi(piProds[sdir][sister]);
							else
								input->addi(hiProds[sdir][sister]);
						}
						else if (sister > i) {
							if (completeInside == CMPL_RIGHT2LEFT)
								input->addi(piProds[sdir][sister]);
							else
								input->addi(hiProds[sdir][sister]);
						}
					}
				}
				input->divi(nSisters)->addi(inputParent);
				tree->outer[i] = func(input);
			}
		}

		// compute partial outers for CHILDREN
		input->fill(0);
		int nLeft = 0;
		//left
		Matrix::gemv(1, CblasNoTrans, params->Woh[0], tree->headInner[i], 0, inputHead);
		Matrix::gemv(1, CblasNoTrans, params->Wop[0], tree->outer[i], 1, inputHead);
		inputHead->addi(params->bo[0]);
		if (tree->nChildren[i][0] == 0) {
			Matrix* temp = inputHead->add(params->anonInner[0]);
			tree->EOCOuter[0][i] = func(temp); delete temp;
			input->copy(params->anonInner[0]);
			nLeft = 1;
		}
		else {
			int leftSister = -1;
			// compute outer for its children
			for (int j = 0; j < tree->nChildren[i][0]; j++) {
				int child = tree->children[i][0][j];
				if (leftSister > -1) {
					input->addi(hiProds[0][leftSister]);
					Matrix* temp = input->div(j)->addi(inputHead);
					tree->pOuter[child] = func(temp); delete temp;
				}
				else {
					Matrix* temp = inputHead->add(params->anonInner[0]);
					tree->pOuter[child] = func(temp); delete temp;
				}

				leftSister = child;
			}

			// compute outer for EOC
			nLeft = tree->nChildren[i][0];
			input->addi(hiProds[0][leftSister]);
			Matrix* temp = input->div(nLeft)->addi(inputHead);
			tree->EOCOuter[0][i] = func(temp); delete temp;
		}

		//right
		Matrix::gemv(1, CblasNoTrans, params->Woh[1], tree->headInner[i], 0, inputHead);
		Matrix::gemv(1, CblasNoTrans, params->Wop[1], tree->outer[i], 1, inputHead);
		inputHead->addi(params->bo[1]);
		if (tree->nChildren[i][1] == 0) {
			Matrix* temp = params->anonInner[1]->add(input)->divi(nLeft+1)->addi(inputHead);
			tree->EOCOuter[1][i] = func(temp); delete temp;
		}
		else {
			int leftSister = -1;

			// compute outer for children
			for (int j = 0; j < tree->nChildren[i][1]; j++) {
				int child = tree->children[i][1][j];
				if (leftSister > -1) {
					input->addi(hiProds[1][leftSister]);
					Matrix* temp = input->div(nLeft+j)->addi(inputHead);
					tree->pOuter[child] = func(temp); delete temp;
				}
				else {
					Matrix* temp = params->anonInner[1]->add(input)->divi(nLeft+1)->addi(inputHead);
					tree->pOuter[child] = func(temp); delete temp;
				}
				leftSister = child;
			}

			input->addi(hiProds[1][leftSister]);
			Matrix* temp = input->div(nLeft+tree->nChildren[i][1])->addi(inputHead);
			tree->EOCOuter[1][i] = func(temp); delete temp;
		}
	}

	delete input;
	delete inputHead;
	delete inputParent;

	// compute softmax
	// Pr(deprel | outer)
	Matrix* score;
	tree->deprelProb = new Matrix*[tree->nNodes];
	tree->deprelProb[0] = Matrix::zeros(params->bdr->length);
	score = new Matrix(params->bdr->length);
	for (int i = 1; i < tree->nNodes; i++) {
		Matrix::gemv(1, CblasNoTrans, params->Wdr, tree->pOuter[i], 0, score);
		score->addi(params->bdr);
		tree->deprelProb[i] = new Matrix(score->length);
		Utils::safelyComputeSoftmax(tree->deprelProb[i]->data, score->data, score->length);
	}

	tree->EOCProb = new Matrix**[2]{ new Matrix*[tree->nNodes], new Matrix*[tree->nNodes]};
	for (int i = 0; i < tree->nNodes; i++)
		for (int dir = 0; dir < 2; dir++) {
			Matrix::gemv(1, CblasNoTrans, params->Wdr, tree->EOCOuter[dir][i], 0, score);
			score->addi(params->bdr);
			tree->EOCProb[dir][i] = new Matrix(score->length);
			Utils::safelyComputeSoftmax(tree->EOCProb[dir][i]->data, score-> data, score->length);
		}
	delete score;

	// Pr(pos | deprel, outer)
	tree->posProb = new Matrix*[tree->nNodes];
	tree->posProb[0] = Matrix::zeros(params->bpos->length);
	score = new Matrix(params->bpos->length);
	for (int i = 1; i < tree->nNodes; i++) {
		Matrix::gemv(1, CblasNoTrans, params->Wpos, tree->pOuter[i], 0, score);
		Matrix* temp = params->Wdrpos->getColumn(tree->deprel[i]);
		score->addi(temp)->addi(params->bpos);
		delete temp;
		tree->posProb[i] = new Matrix(score->length);
		Utils::safelyComputeSoftmax(tree->posProb[i]->data, score->data, score->length);
	}
	delete score;

	// Pr(word | ...)
	tree->wordProb = new Matrix*[tree->nNodes];
	tree->wordProb[0] = Matrix::zeros(1); // undefined
	for (int i = 1; i < tree->nNodes; i++) {
		int word = tree->word[i];
		int* path = vocaDic->path[word];
		real* code = vocaDic->code[word];
		int codeLen = vocaDic->codeLength[word];
		tree->wordProb[i] = Matrix::zeros(codeLen);
		for (int j = 0; j < codeLen; j++) {
			tree->wordProb[i]->data[j] = Matrix::dot(params->Wword[path[j]], tree->pOuter[i]);
		}
		Matrix* temp1 = params->Wdrword->getRows(path, codeLen);
		Matrix* temp2 = temp1->getColumn(tree->deprel[i]);
		tree->wordProb[i]->addi(temp2);
		delete temp1; delete temp2;

		temp1 = params->Wposword->getRows(path, codeLen);
		temp2 = temp1->getColumn(tree->pos[i]);
		tree->wordProb[i]->addi(temp2);
		delete temp1; delete temp2;

		temp1 = params->bword->getRows(path, codeLen);
		temp2 = new Matrix(codeLen, code);
		tree->wordProb[i]->addi(temp1)->muli(temp2);
		delete temp1; delete temp2;

		temp1 = tree->wordProb[i];
		tree->wordProb[i] = sigmoid(temp1);
		delete temp1;
	}

	// Pr(cap | ...)
	tree->capProb = new Matrix*[tree->nNodes];
	tree->capProb[0] = Matrix::zeros(params->bcap->length);
	score = new Matrix(params->bcap->length);
	for (int i = 1; i < tree->nNodes; i++) {
		Matrix::gemv(1, CblasNoTrans, params->Wcap, tree->pOuter[i], 0, score);
		Matrix* temp1 = params->Wdrcap->getColumn(tree->deprel[i]);
		Matrix* temp2 = params->Wposcap->getColumn(tree->pos[i]);
		Matrix* temp3 = params->Wwordcap->getColumn(tree->word[i]);
		score->addi(temp1)->addi(temp2)->addi(temp3)->addi(params->bcap);
		delete temp1; delete temp2; delete temp3;

		tree->capProb[i] = new Matrix(score->length);
		Utils::safelyComputeSoftmax(tree->capProb[i]->data, score->data, score->length);
	}
	delete score;

	// compute errors
	real totalError = 0;
	for (int i = 1; i < tree->nNodes; i++)
		totalError -= log(tree->deprelProb[i]->data[tree->deprel[i]])
		+ log(tree->EOCProb[0][i]->data[deprelDic->size()])
		+ log(tree->EOCProb[1][i]->data[deprelDic->size()])
		+ log(tree->posProb[i]->data[tree->pos[i]])
		+ Utils::sumLog(tree->wordProb[i]->data, tree->wordProb[i]->length)
	+ log(tree->capProb[i]->data[tree->cap[i]]);

	totalError -= log(tree->EOCProb[0][0]->data[deprelDic->size()])
											+ log(tree->EOCProb[1][0]->data[deprelDic->size()]);

	for (int i = 1; i < tree->nNodes; i++) {
		int dir = tree->dir[i];
		delete hiProds[dir][i];
		if (completeInside != CMPL_NONE)
			delete piProds[dir][i];
	}
	delete[] piProds[0];
	delete[] piProds[1];
	delete[] piProds;
	delete[] hiProds[0];
	delete[] hiProds[1];
	delete[] hiProds;


	return totalError;
}

void IORNN::backpropOutside(Deptree* tree, Param* grad, int cmplInside) {
	tree->gradheadI = new Matrix*[tree->nNodes];
	tree->gradphraseI = new Matrix*[tree->nNodes];
	tree->gradO = new Matrix*[tree->nNodes];
	tree->gradpO = new Matrix*[tree->nNodes];
	tree->gradEOCO = new Matrix**[2] { new Matrix*[tree->nNodes],
		new Matrix*[tree->nNodes] };

	for (int i = 0; i < tree->nNodes; i++) {
		tree->gradheadI[i] = Matrix::zeros(dim);
		tree->gradphraseI[i] = Matrix::zeros(dim);
		tree->gradO[i] = Matrix::zeros(dim);
		tree->gradpO[i] = Matrix::zeros(dim);
		tree->gradEOCO[0][i] = Matrix::zeros(dim);
		tree->gradEOCO[1][i] = Matrix::zeros(dim);
	}

	tree->gZdr = tree->deprelProb;
	tree->gZpos = tree->posProb;
	tree->gZword = tree->wordProb;
	tree->gZcap = tree->capProb;
	tree->gZ_EOC = tree->EOCProb;

	for (int i = 1; i < tree->nNodes; i++) {
		tree->gZdr[i]->data[tree->deprel[i]]--;
		tree->gZpos[i]->data[tree->pos[i]]--;
		tree->gZcap[i]->data[tree->cap[i]]--;
		Matrix* temp = new Matrix(vocaDic->codeLength[tree->word[i]], vocaDic->code[tree->word[i]]);
		tree->gZword[i]->addi(-1)->muli(temp); delete temp;
	}

	tree->gZdr[0] = Matrix::zeros(tree->gZdr[1]->length);
	tree->gZpos[0] = Matrix::zeros(tree->gZpos[1]->length);
	tree->gZcap[0] = Matrix::zeros(tree->gZcap[1]->length);
	for (int i = 0; i < tree->nNodes; i++) {
		tree->gZ_EOC[0][i]->data[deprelDic->size()]--;
		tree->gZ_EOC[1][i]->data[deprelDic->size()]--;
	}

	// for Pr (. | context)
	for (int i = 0; i < tree->nNodes; i++) {
		Matrix::ger(1, tree->gZdr[i], tree->pOuter[i], grad->Wdr);
		Matrix::ger(1, tree->gZ_EOC[0][i], tree->EOCOuter[0][i], grad->Wdr);
		Matrix::ger(1, tree->gZ_EOC[1][i], tree->EOCOuter[1][i], grad->Wdr);
		grad->bdr->addi(tree->gZdr[i])->addi(tree->gZ_EOC[0][i])->addi(tree->gZ_EOC[1][i]);

		Matrix::gemv(1, CblasTrans, params->Wdr, tree->gZdr[i], 1, tree->gradpO[i]);
		Matrix::gemv(1, CblasTrans, params->Wdr, tree->gZ_EOC[0][i], 1, tree->gradEOCO[0][i]);
		Matrix::gemv(1, CblasTrans, params->Wdr, tree->gZ_EOC[1][i], 1, tree->gradEOCO[1][i]);

		Matrix::ger(1, tree->gZpos[i], tree->pOuter[i], grad->Wpos);
		grad->bpos->addi(tree->gZpos[i]);
		Matrix::gemv(1, CblasTrans, params->Wpos, tree->gZpos[i], 1, tree->gradpO[i]);

		Matrix::ger(1, tree->gZcap[i], tree->pOuter[i], grad->Wcap);
		grad->bcap->addi(tree->gZcap[i]);
		Matrix::gemv(1, CblasTrans, params->Wcap, tree->gZcap[i], 1, tree->gradpO[i]);

	}

	for (int i = 1; i < tree->nNodes; i++) {
		int word = tree->word[i];
		int* path = vocaDic->path[word];
		int codeLen = vocaDic->codeLength[word];
		int dr = tree->deprel[i];
		int pos = tree->pos[i];

		grad->Wdrpos->addColumn(dr, tree->gZpos[i]);
		for (int d = 0; d < dim; d++)
			for (int j = 0; j < codeLen; j++)
				tree->gradpO[i]->data[d] += params->Wword[path[j]]->data[d] * tree->gZword[i]->data[j];

		for (int j = 0; j < codeLen; j++) {
			int p = path[j];
			real g = tree->gZword[i]->data[j];

			Matrix::axpy(g, tree->pOuter[i], grad->Wword[p]);
			grad->bword->data[p] += g;
			grad->Wdrword->put(p, dr, grad->Wdrword->get(p, dr) + g);
			grad->Wposword->put(p, pos, grad->Wposword->get(p, pos) + g);
		}

		grad->Wdrcap->addColumn(dr, tree->gZcap[i]);
		grad->Wposcap->addColumn(pos, tree->gZcap[i]);
		grad->Wwordcap->addColumn(word, tree->gZcap[i]);
	}

	// backward
	Matrix** gradZpO = new Matrix*[tree->nNodes];
	Matrix*** gradZEOCO = new Matrix**[2] { new Matrix*[tree->nNodes], new Matrix*[tree->nNodes] };

	for (int i = 0; i < tree->nNodes; i++) {
		Matrix *temp = funcPrime(tree->pOuter[i]);
		gradZpO[i] = tree->gradpO[i]->muli(temp); delete temp;
		for (int dir = 0; dir < 2; dir++) {
			temp = funcPrime(tree->EOCOuter[dir][i]);
			gradZEOCO[dir][i] = tree->gradEOCO[dir][i]->muli(temp); delete temp;
			Matrix::ger(1, gradZEOCO[dir][i], tree->headInner[i], grad->Woh[dir]);
			Matrix::ger(1, gradZEOCO[dir][i], tree->outer[i], grad->Wop[dir]);
			grad->bo[dir]->addi(gradZEOCO[dir][i]);

			Matrix::gemv(1, CblasTrans, params->Woh[dir], gradZEOCO[dir][i], 1, tree->gradheadI[i]);
			Matrix::gemv(1, CblasTrans, params->Wop[dir], gradZEOCO[dir][i], 1, tree->gradO[i]);
		}
	}

	Matrix*** phraseISums = new Matrix**[2] {new Matrix*[tree->nNodes], new Matrix*[tree->nNodes]};
	Matrix*** sums = new Matrix**[2] {new Matrix*[tree->nNodes], new Matrix*[tree->nNodes]};
	for (int i = 0; i < tree->nNodes; i++)
		for (int dir = 0; dir < 2; dir++) {
			phraseISums[dir][i] = Matrix::zeros(dim);
			sums[dir][i] = Matrix::zeros(dim);
		}

	// for partial outer
	for (int i = tree->nNodes-1; i >= 0; i--) {
		// for EOC
		// left
		Matrix* gz = gradZEOCO[0][i];
		int nLeft = 1;

		if (tree->nChildren[i][0] == 0)
			grad->anonInner[0]->addi(gz);
		else {
			nLeft = tree->nChildren[i][0];
			real t = (real)1. / nLeft;
			for (int j = 0; j < tree->nChildren[i][0]; j++) {
				int child = tree->children[i][0][j];
				Matrix::axpy(t, gz, sums[0][child]);
			}
		}

		// right
		gz = gradZEOCO[1][i];
		int nRight = 1;

		if (tree->nChildren[i][1] == 0)
			grad->anonInner[1]->addi(1./(nLeft+1), gz);
		else {
			nRight = tree->nChildren[i][1];
			real t = 1.f / (nRight + nLeft);
			for (int j = 0; j < tree->nChildren[i][1]; j++) {
				int child = tree->children[i][1][j];
				Matrix::axpy(t, gz, sums[1][child]);
			}
		}

		if (tree->nChildren[i][0] == 0)
			grad->anonInner[0]->addi((real)1./(nLeft+nRight), gz);
		else {
			real t = (real)1. / (nRight + nLeft);
			for (int j = 0; j < tree->nChildren[i][0]; j++) {
				int child = tree->children[i][0][j];
				Matrix::axpy(t, gz, sums[0][child]);
			}
		}

		// for children's partial outer
		// left
		for (int j = 0; j < tree->nChildren[i][0]; j++) {
			int child = tree->children[i][0][j];
			gz = gradZpO[child];

			Matrix::ger(1, gz, tree->headInner[i], grad->Woh[0]);
			Matrix::ger(1, gz, tree->outer[i], grad->Wop[0]);
			grad->bo[0]->addi(gz);

			Matrix::gemv(1, CblasTrans, params->Woh[0], gz, 1, tree->gradheadI[i]);
			Matrix::gemv(1, CblasTrans, params->Wop[0], gz, 1, tree->gradO[i]);

			if (j == 0)
				grad->anonInner[0]->addi(gz);
			else {
				real t = (real)1. / j;
				for (int k = 0; k < j; k++) {
					int sister = tree->children[i][0][k];
					Matrix::axpy(t, gz, sums[0][sister]);
				}
			}
		}

		// right
		for (int j = 0; j < tree->nChildren[i][1]; j++) {
			int child = tree->children[i][1][j];
			gz = gradZpO[child];

			Matrix::ger(1, gz, tree->headInner[i], grad->Woh[1]);
			Matrix::ger(1, gz, tree->outer[i], grad->Wop[1]);
			grad->bo[1]->addi(gz);

			Matrix::gemv(1, CblasTrans, params->Woh[1], gz, 1, tree->gradheadI[i]);
			Matrix::gemv(1, CblasTrans, params->Wop[1], gz, 1, tree->gradO[i]);

			nRight = 1;
			if (j == 0)
				grad->anonInner[1]->addi((real)1./(nLeft+1), gz);
			else {
				nRight = j;
				real t = (real)1. / (nLeft + nRight);
				for (int k = 0; k < j; k++) {
					int sister = tree->children[i][1][k];
					Matrix::axpy(t, gz, sums[1][sister]);
				}
			}

			if (tree->nChildren[i][0] == 0)
				grad->anonInner[0]->addi((real)1./(nLeft+nRight), gz);
			else {
				real t = (real)1. / (nLeft + nRight);
				for (int k = 0; k < tree->nChildren[i][0]; k++) {
					int sister = tree->children[i][0][k];
					Matrix::axpy(t, gz, sums[0][sister]);
				}
			}
		}

		// for full outer
		if (i == 0)
			grad->anonOuter->addi(tree->gradO[0]);
		else {
			int parent = tree->parent[i];
			int dir = tree->dir[i];
			Matrix* temp = funcPrime(tree->outer[i]);
			gz = tree->gradO[i]->muli(temp); delete temp;

			Matrix::ger(1, gz, tree->headInner[parent], grad->Woh[dir]);
			Matrix::ger(1, gz, tree->outer[parent], grad->Wop[dir]);
			grad->bo[dir]->addi(gz);

			Matrix::gemv(1, CblasTrans, params->Woh[dir], gz, 1, tree->gradheadI[parent]);
			Matrix::gemv(1, CblasTrans, params->Wop[dir], gz, 1, tree->gradO[parent]);

			int nSisters = tree->nChildren[parent][0] + tree->nChildren[parent][1] - 1;
			if (nSisters == 0)
				grad->anonInner[dir]->addi(gz);
			else {
				real t = (real)1. / nSisters;
				for (int cdir = 0; cdir < 2; cdir++) {
					for (int j = 0; j < tree->nChildren[parent][cdir]; j++) {
						int sister = tree->children[parent][cdir][j];
						if (sister < i) {
							if (cmplInside == CMPL_LEFT2RIGHT)
								Matrix::axpy(t, gz, phraseISums[cdir][sister]);
							else
								Matrix::axpy(t, gz, sums[cdir][sister]);
						}
						else if (sister > i) {
							if (cmplInside == CMPL_RIGHT2LEFT)
								Matrix::axpy(t, gz, phraseISums[cdir][sister]);
							else
								Matrix::axpy(t, gz, sums[cdir][sister]);
						}
					}
				}
			}
		}
	}

	for (int i = 1; i < tree->nNodes; i++) {
		int dir = tree->dir[i];
		int dr = tree->deprel[i];
		if (cmplInside != CMPL_NONE) {
			Matrix::ger(1, phraseISums[dir][i], tree->phraseInner[i], grad->Wo[dir][dr]);
			Matrix::gemv(1, CblasTrans, params->Wo[dir][dr], phraseISums[dir][i], 1, tree->gradphraseI[i]);
		}

		Matrix::ger(1, sums[dir][i], tree->headInner[i], grad->Wo[dir][dr]);
		Matrix::gemv(1, CblasTrans, params->Wo[dir][dr], sums[dir][i], 1, tree->gradheadI[i]);
	}

	Matrix::free(phraseISums, 2, tree->nNodes);
	Matrix::free(sums, 2, tree->nNodes);

	delete[] gradZpO;
	delete[] gradZEOCO[0];
	delete[] gradZEOCO[1];
	delete[] gradZEOCO;
}

void IORNN::backpropInside(Deptree* tree, Param* grad, int cmplInside) {

	// for head inner
	grad->rootInner->addi(tree->gradheadI[0]);
	Matrix** gradZheadI = new Matrix*[tree->nNodes];
	for (int i = 0; i < tree->nNodes; i++) {
		Matrix* temp = funcPrime(tree->headInner[i]);
		gradZheadI[i] = tree->gradheadI[i]->muli(temp); delete temp;
	}

	for (int i = 1; i < tree->nNodes; i++) {
		int word = tree->word[i];
		int pos = tree->pos[i];
		int cap = tree->cap[i];

		Matrix* gz = gradZheadI[i];
		Matrix::ger(1, gz, params->L[word], grad->Ww);
		grad->bw->addi(gz);
		Matrix::gemv(1, CblasTrans, params->Ww, gz, 1, grad->L[word]);
		grad->Lpos[pos]->addi(gz);
		grad->Lcap[cap]->addi(gz);
	}
	delete[] gradZheadI;

	// for phrase inners
	if (cmplInside != CMPL_NONE) {
		for (int i = 0; i < tree->nNodes; i++) {
			int word = tree->word[i];
			int pos = tree->pos[i];
			int cap = tree->cap[i];

			int nChildren = tree->nChildren[i][0] + tree->nChildren[i][1];
			Matrix * temp = funcPrime(tree->phraseInner[i]);
			Matrix* gz = tree->gradphraseI[i]->muli(temp); delete temp;

			// for head
			Matrix* gradZhead = gz->dup();
			if (nChildren > 0) {
				Matrix::gemv(1, CblasTrans, params->Wih, gz, 0, gradZhead);
				Matrix *temp = funcPrime(tree->headInner[i]);
				gradZhead->muli(temp); delete temp;
				Matrix::ger(1, gz, tree->headInner[i], grad->Wih);
			}

			if (i == 0) //root
					Matrix::gemv(1, CblasTrans, params->Wih, gz, 1, grad->rootInner);
			else {
				Matrix::ger(1, gradZhead, params->L[word], grad->Ww);
				grad->bw->addi(gradZhead);
				Matrix::gemv(1, CblasTrans, params->Ww, gradZhead, 1, grad->L[word]);
				grad->Lpos[pos]->addi(gradZhead);
				grad->Lcap[cap]->addi(gradZhead);
			}
			delete gradZhead;


			// for children
			if (nChildren > 0) {
				grad->bi->addi(gz);
				gz->divi(nChildren);
				for (int dir = 0; dir < 2; dir++) {
					for (int j = 0; j < tree->nChildren[i][dir]; j++) {
						int child = tree->children[i][dir][j];
						int dr = tree->deprel[child];

						Matrix::ger(1, gz, tree->phraseInner[child], grad->Wi[dir][dr]);
						Matrix::gemv(1, CblasTrans, params->Wi[dir][dr], gz, 1, tree->gradphraseI[child]);
					}
				}
			}
		}
	}
}

vector<real> IORNN::computeLogProb(vector<Deptree*>& trees) {
	vector<real> scores;

	for (int i = 0; i < trees.size(); i++) {
		forwardInside(trees[i]);
		real lp = -forwardOutside(trees[i]);
		scores.push_back(lp);
		trees[i]->freeTempArrays();
	}
	return scores;
}

void** IORNN::computeCostAndGrad(Treebank* tb, int startId, int endId, Param** gradList) {

	unordered_set<int>* tWords = new unordered_set<int>;
	unordered_set<int>* tPath = new unordered_set<int>;
	int nSample = 0;

	// prepare data, the first entry is the main
	for (int i = 0; i < nThreads; i++)
		gradList[i]->fill(0);

	real* costList = new real[nThreads]();
	int step = ceil((endId - startId+1) / (float)nThreads);

#pragma omp parallel for num_threads(nThreads)
	for (int th = 0; th < nThreads; th++) {
		for (int i = startId + th*step; i <= min(endId, startId + (th+1)*step-1); i++) {
			Deptree* tree = tb->get(i);
			forwardInside(tree);
			costList[th] += forwardOutside(tree);

			backpropOutside(tree, gradList[th], completeInside);
			backpropInside(tree, gradList[th], completeInside);
			tree->freeTempArrays();
		}
	}

	for (int i = startId; i <= endId; i++) {
		Deptree* tree = tb->get(i);
		nSample += tree->nNodes;
		for (int k = 1; k < tree->nNodes; k++) {
			int word = tree->word[k];
			tWords->insert(word);
			int* path = vocaDic->path[word];
			int codeLen = vocaDic->codeLength[word];
			for (int p = 0; p < codeLen; p++)
				tPath->insert(path[p]);
		}
	}

	Param* grad = gradList[0];

	// merge
	for (int i = 1; i < nThreads; i++) {
		for (int j = 0; j < grad->weights.size(); j++)
			grad->weights[j]->addi(gradList[i]->weights[j]);

		costList[0] += costList[i];
	}

	// compute regularities
	real cost = costList[0];

	cost /= nSample;
	for (int i = 0; i < params->nMatricesWoWE; i++) {
		Matrix* W = params->weights[i];
		cost += Utils::sumSqr(W->data, W->length) * lambda/2;
		grad->weights[i]->divi(nSample)->addi(lambda, W);
	}

	for (unordered_set<int>::iterator it = tWords->begin(); it != tWords->end(); it++) {
		int word = *it;
		Matrix* wemb = params->L[word];
		cost += Utils::sumSqr(wemb->data, wemb->length) * lambdaL/2;
		grad->L[word]->divi(nSample)->addi(lambdaL, wemb);
	}

	for (unordered_set<int>::iterator it = tPath->begin(); it != tPath->end(); it++) {
		int p = *it;
		Matrix* v = params->Wword[p];
		cost += Utils::sumSqr(v->data, v->length) * lambda/2;
		grad->Wword[p]->divi(nSample)->addi(lambda, v);
	}

	delete[] costList;

	return new void*[3] {
		&cost,
		tWords,
		tPath };
}


// make sure gradients are computed correctly
void IORNN::checkGradient(Treebank* tb) {
	real epsilon = 1e-4f;
	Param** grad = new Param*[nThreads];
	for (int i = 0; i < nThreads; i++)
		grad[i] = createGrad();
	void** output = computeCostAndGrad(tb, 0, tb->size()-1, grad);

	vector<real> gradV;
	vector<real> numGradV;
	Param** otherGrad = new Param*[nThreads];
	for (int i = 0; i < nThreads; i++)
		otherGrad[i] = createGrad();

	for (int i = 0; i < params->weights.size(); i++) {
		Matrix* W = params->weights[i];
		for (int r = 0; r < W->rows; r++)
			for (int c = 0; c < W->cols; c++) {
				gradV.push_back(grad[0]->weights[i]->get(r,c));

				// plus
				W->put(r, c, W->get(r,c)+epsilon);
				output = computeCostAndGrad(tb, 0, tb->size()-1, otherGrad);
				real pluscost = *(real*)output[0];

				// minus
				W->put(r, c, W->get(r,c)-2*epsilon);
				output = computeCostAndGrad(tb, 0, tb->size()-1, otherGrad);
				real minuscost = *(real*)output[0];

				W->put(r, c, W->get(r,c)+epsilon);

				numGradV.push_back((pluscost - minuscost) / (2*epsilon));
				real diff = abs(numGradV.back() - gradV.back());
				if (diff > 0 && abs(gradV.back()) < 1e-7)
					cout << diff << " " << i << " " << r << " " << c << " : " << numGradV.back() << " " << gradV.back() << endl;
				if (diff > 5e-9)
					cout << diff << " " << i << " " << r << " " << c << " : " << numGradV.back() << " " << gradV.back() << endl;
			}
	}

	Matrix* X = new Matrix((int)gradV.size(), &gradV[0]);
	Matrix* Y = new Matrix((int)numGradV.size(), &numGradV[0]);
	real diff = Matrix::nrm2(X->add(-1, Y)) / Matrix::nrm2(X->add(Y));
	cout << diff << endl;
	cout << "should be < 1e-9" << endl;

	for (int i = 0; i < nThreads; i++) {
		Param::free(grad[i]);
		Param::free(otherGrad[i]);
	}
	delete[] grad;
	delete[] otherGrad;
	delete[] output;
	delete X;
	delete Y;
}
